package com.jqj.example;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.alibaba.ttl.threadpool.TtlExecutors;
import jakarta.validation.constraints.NotNull;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.io.IOException;
import java.time.LocalTime;
import java.util.concurrent.*;

/*
 * @Desc:
 * @Author: lijingzheng
 * @Date: 2024/11/27
 */
public class ThreadLocalTest {
//    -javaagent:C:\Users\EDY\IdeaProjects\MyJavaAgentByMaven\target\MyJavaAgentByMaven-1.0-SNAPSHOT.jar

//    private static final InheritableThreadLocal<Integer> THREAD_LOCAL = new InheritableThreadLocal<>();

    private static final TransmittableThreadLocal<Integer> THREAD_LOCAL = new TransmittableThreadLocal<>(true);

    public static void main1(String[] args) throws IOException {
        OauthContext.set(123);
        System.out.printf("\033[33;1m%s ===>> %s ===>> %s\033[0m\n",
                LocalTime.now(), Thread.currentThread(), OauthContext.get());

        ThreadPoolTaskExecutor pool = taskExecutor();
        for (int i = 0; i < 12; i++) {
            int finalI = i;
            CompletableFuture.runAsync(() -> {
                try {
                    TimeUnit.MILLISECONDS.sleep(5000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
                System.out.printf("\033[32;1m%s ===>> %s ===>> %s ===>> %s ===>> QueueSize=%s, ActiveCount=%s\033[0m\n",
                        LocalTime.now(), finalI, Thread.currentThread(), OauthContext.get(), pool.getThreadPoolExecutor().getQueue().size(), pool.getActiveCount());
            }, pool);
            System.out.printf("\033[34;1m%s ===>> %s ===>> QueueSize=%s, ActiveCount=%s\033[0m\n", LocalTime.now(),
                    finalI, pool.getThreadPoolExecutor().getQueue().size(), pool.getActiveCount());
        }

        OauthContext.set(234);
        try {
            TimeUnit.MILLISECONDS.sleep(5000);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }

        for (int i = 12; i < 20; i++) {
            int finalI = i;
            CompletableFuture.runAsync(() -> {
                try {
                    TimeUnit.MILLISECONDS.sleep(5000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
                System.out.printf("\033[32;1m%s ===>> %s ===>> %s ===>> %s ===>> QueueSize=%s, ActiveCount=%s\033[0m\n",
                        LocalTime.now(), finalI, Thread.currentThread(), OauthContext.get(), pool.getThreadPoolExecutor().getQueue().size(), pool.getActiveCount());
            }, pool);
            System.out.printf("\033[34;1m%s ===>> %s ===>> QueueSize=%s, ActiveCount=%s\033[0m\n", LocalTime.now(),
                    finalI, pool.getThreadPoolExecutor().getQueue().size(), pool.getActiveCount());

        }

        System.in.read();
    }

    public static void main(String[] args) throws Exception {
        THREAD_LOCAL.set(123);

        Executor pool = TtlExecutors.getTtlExecutor(taskExecutor());
        for (int i = 0; i < 10; i++) {
            int finalI = i;
            CompletableFuture.runAsync(() -> {
                try {
                    TimeUnit.MILLISECONDS.sleep(5000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
                System.out.printf("\033[32;1m%s ===>> %s ===>> %s ===>> %s\033[0m\n",
                        LocalTime.now(), finalI, Thread.currentThread(), THREAD_LOCAL.get());
//                THREAD_LOCAL.set(345);
//                THREAD_LOCAL.remove();
            }, pool);
            System.out.printf("\033[34;1m%s ===>> %s\033[0m\n", LocalTime.now(), finalI);
        }

        THREAD_LOCAL.set(234);
        try {
            TimeUnit.SECONDS.sleep(15);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }

        for (int i = 12; i < 30; i++) {
            int finalI = i;
            CompletableFuture.runAsync(() -> {
                try {
                    TimeUnit.MILLISECONDS.sleep(5000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
                System.out.printf("\033[36;1m%s ===>> %s ===>> %s ===>> %s\033[0m\n",
                        LocalTime.now(), finalI, Thread.currentThread(), THREAD_LOCAL.get());
            }, pool);
            System.out.printf("\033[34;1m%s ===>> %s\033[0m\n", LocalTime.now(), finalI);
        }

        System.in.read();
    }

    public static ThreadPoolTaskExecutor taskExecutor() {
        ThreadPoolTaskExecutor poolTaskExecutor = new ThreadPoolTaskExecutor();
        poolTaskExecutor.setCorePoolSize(3);
        poolTaskExecutor.setMaxPoolSize(6);
        // 设置队列容量
        poolTaskExecutor.setQueueCapacity(5);
        // 设置线程活跃时间（秒）
        poolTaskExecutor.setKeepAliveSeconds(2);
        //设置TaskDecorator，用于解决父子线程间的数据复用
        poolTaskExecutor.setTaskDecorator(new ContextTaskDecorator());
        poolTaskExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.DiscardPolicy());
        // 等待所有任务结束后再关闭线程池
        poolTaskExecutor.setWaitForTasksToCompleteOnShutdown(true);
        poolTaskExecutor.initialize();
        return poolTaskExecutor;
    }

    public static class ContextTaskDecorator implements TaskDecorator {

        @NotNull
        @Override
        public Runnable decorate(@NotNull Runnable runnable) {
            //获取父线程的loginVal
            Integer loginVal = OauthContext.get();
            return () -> {
                try {
                    // 将主线程的请求信息，设置到子线程中
                    OauthContext.set(loginVal);
//                    System.out.printf("\033[35;1m%s ===>> %s ===>> %s\033[0m\n", LocalTime.now(),
//                            Thread.currentThread(), OauthContext.get());
                    // 执行子线程，这一步不要忘了
                    runnable.run();
                } finally {
                    // 线程结束，清空这些信息，否则可能造成内存泄漏
                    OauthContext.clear();
                }
            };
        }
    }

    public static class OauthContext {

        private static final ThreadLocal<Integer> loginValThreadLocal = new ThreadLocal<>();

        public static Integer get() {
            return loginValThreadLocal.get();
        }

        public static void set(Integer loginVal) {
            loginValThreadLocal.set(loginVal);
        }

        public static void clear() {
            loginValThreadLocal.remove();
        }
    }

}
